# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import functools
import math
import logging
import signal

from absl import (flags, app)
from tornado.ioloop import IOLoop, PeriodicCallback
import numpy as np

import coin.flow.subscriber as flow_subscriber
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.base.order_gateway import OrderSide
from coin.base.book.types import TradeSide
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.strategy.hk_hard_arb.products import NativeProducts
from coin.base.param_util import to_list
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_order_checker_pb2 import AccountCheckerConfig
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)
from coin.strategy.accounting.pnl_balance_printer.util import reserve_as_proto
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.passive.util import \
  convert_executor_config_from_dict as convert_pass_config
from coin.support.accounting.checker.account_checker import AccountChecker
from coin.support.accounting.checker.config_util import convert_checker_config_from_dict
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.exchange.huobi_futures.kr_rest.native_private_client import HuobiFuturesOrderSide
from coin.exchange.okex_v3.kr_rest.native_private_client import OkexFuturesOrderSide
from coin.strategy.tool.fill_log_dumper import FillLogDumper
from coin.strategy.basis_strat.fill_info_writer import FillInfo, FillInfoWriter
from coin.strategy.util.request_util import convert_strategy_request_from_dict
from coin.strategy.accounting.fee import get_fee_rate
from coin.strategy.basis_strat.moving_average import MovingAverage
from coin.support.telemetry.latency_recorder import LatencyRecorder, LifeOfSignal
from coin.strategy.basis_strat.base_strategy import BaseStrategy

FLAGS = flags.FLAGS


class DummyPosition(object):
  def __init__(self):
    self.long = self.short = 0


def _gen_arb_products(common, products0, products1, ProductType0, ProductType1):
  if 'exch0_quote_currency' in common and 'exch1_quote_currency' in common:
    exch0_quote_currency = common['exch0_quote_currency']
    exch1_quote_currency = common['exch1_quote_currency']
  else:
    raise ValueError(common)

  products0 = NativeProducts.FromTypedStrProduct(exch0_quote_currency,
                                                 sorted(list(products0.keys())),
                                                 ProductType0)
  products1 = NativeProducts.FromTypedStrProduct(exch1_quote_currency,
                                                 sorted(list(products1.keys())),
                                                 ProductType1)
  assert len(products0.products) == len(products1.products)
  return products0, products1


class BasisPassSubStrategy(object):
  def __init__(self,
               parent,
               product0,
               product1,
               fillinfo_writer,
               strategy_logger,
               latency_recorder=None):
    self._parent = parent
    self._product0 = product0
    self._product1 = product1
    self._fillinfo_writer = fillinfo_writer
    self._strategy_logger = strategy_logger
    self._latency_recorder = latency_recorder
    self._ticker0 = product0.full_symbol
    self._ticker1 = product1.full_symbol
    self._tick_info1 = (get_holder_from_product(product1).product_info.gen_tick(),
                        get_holder_from_product(product1).product_info.gen_tick_qty())
    if hasattr(self._tick_info1[0], 'tick_size'):
      ticksize = self._tick_info1[0].tick_size
      self._next_price_func = lambda p: p + ticksize - (p % ticksize)
      self._prev_price_func = lambda p: p - ((p % ticksize) or ticksize)
    else:
      self._next_price_func = self._tick_info1[0].get_next_price
      self._prev_price_func = self._tick_info1[0].get_prev_price

    self.og1 = None
    self._pass_executor = None
    self.logger = parent.logger
    self._last_feed = parent._last_feed

    self._equal_theshold = None
    self._lot_size = None
    self._max_pos = None
    self._min_pos = None
    self._agg_taker_fee = None
    self._last_submit_fill_ts = 0
    self._last_book = {p: None for p in [product0, product1]}

    # Decide if book has changed.
    self._book_change_levels = min(FLAGS.book_change_levels, FLAGS.book_levels)

    self._get_pos_func = None
    self._is_buy_order_func = None
    self._dummy_position = DummyPosition()
    self._raw_run_size = 1.0
    self._relative_run_size = None

    self._ma = MovingAverage(FLAGS.window_period_sec * 1e9)
    self._buy_st = 0
    self._sell_st = 0
    self._acc_buy_qty = 0
    self._acc_sell_qty = 0
    self._lambda = FLAGS.discount_factor
    assert self._lambda <= 0
    self._last_processed_trade_ts = 0
    self._buy_restrict_by_basis = False
    self._sell_restrict_by_basis = False
    self._buy_restrict_by_book = False
    self._sell_restrict_by_book = False
    self._buy_restrict_by_trade = False
    self._sell_restrict_by_trade = False
    self._trade_restrict_threshold = FLAGS.trade_imbalance_restrict_threshold
    self._book_restrict_threshold = FLAGS.book_imbalance_restrict_threshold
    self._basis_restrict_threshold = FLAGS.basis_restrict_threshold

  def ready(self):
    # 30s warm up time.
    return self.get_ts() - self._parent._start_time > 30e9

  def get_mid_price(self, bids, asks):
    return (bids[0][0] + asks[0][0]) / 2.0

  def _init_get_pos_func(self, og):
    product_type = self._product1.product_type
    exchange = self._product1.exchange
    if product_type == 'Futures':
      if exchange == 'Huobi':
        self._get_pos_func = lambda: \
            (og.get_huobi_futures_position(self._product1) or self._dummy_position).long - \
            (og.get_huobi_futures_position(self._product1) or self._dummy_position).short
        self._is_buy_order_func = lambda order_side: order_side in [
            HuobiFuturesOrderSide.BUY_OPEN, HuobiFuturesOrderSide.BUY_CLOSE, OrderSide.BUY
        ]
      elif exchange == 'Kraken':
        self._get_pos_func = lambda: \
            getattr(og.og_info.get_product_position_proto(self._product1), 'net_position', 0.0)
        self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
      elif exchange == 'Okex':
        self._get_pos_func = lambda: \
            (og.get_okex_futures_position(self._product1) or self._dummy_position).long - \
            (og.get_okex_futures_position(self._product1) or self._dummy_position).short
        self._is_buy_order_func = lambda order_side: order_side in \
            [OkexFuturesOrderSide.BUY_OPEN, OkexFuturesOrderSide.BUY_CLOSE]
      elif exchange == 'Bitmex':
        self._get_pos_func = lambda: og.get_position(self._product1)
        self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
    elif product_type == 'CurrencyPair':
      self._get_pos_func = lambda: og.get_balance(self._product1.base.currency)
      self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
    assert self._get_pos_func is not None, '%s %s' % (product_type, exchange)

  def write_executor_config(self):
    if self._config1_proto is not None:
      self._strategy_logger.write_executor_config(self._config1_proto)

  def init_executors(self, req1, config1, og1):
    self._relative_run_size = 1.0
    self.og1 = og1
    self._raw_config1 = config1.copy()
    self._config1 = config1.copy()
    self._config1['force_fire'] = True

    if FLAGS.run_size != 1.0:
      assert 0 < FLAGS.run_size < 1, FLAGS.run_size
      self._raw_run_size = FLAGS.run_size
      max_pos = self._config1['max_pos']
      min_pos = self._config1['min_pos']
      neutral_pos = (max_pos + min_pos) / 2.0
      max_exposure = (max_pos - min_pos) / 2.0 * FLAGS.run_size
      self._config1['max_pos'] = neutral_pos + max_exposure
      self._config1['min_pos'] = neutral_pos - max_exposure
      self._config1['lot_size'] = self._config1['lot_size'] * FLAGS.run_size
      self.logger.info('Trimmed size %f %f' % (self._config1['min_pos'], self._config1['max_pos']))

    self._config1_proto = convert_pass_config(self._product1.symbol, self._config1)
    self._pass_executor = construct_mm_executor(req1,
                                                self._config1_proto,
                                                self.og1,
                                                self.logger,
                                                latency_recorder=self._latency_recorder)

    self._init_get_pos_func(og1)
    self._equal_theshold = max(self._config1['lot_size'] * 1e-4, 1e-5)
    self._lot_size = self._config1['lot_size']
    self._max_pos = self._config1['max_pos']
    self._min_pos = self._config1['min_pos']

    _, self._agg_taker_fee = get_fee_rate(self.og1.market_type, self.og1.exchange, self._product1)
    self._raw_agg_taker_fee = self._agg_taker_fee
    if hasattr(self._product1, 'expiry') and self._product1.expiry:
      self.logger.info('Futures product, set fees to 0, %s' % self._product1.expiry)
      self._agg_taker_fee = 0.0
    assert abs(self._agg_taker_fee) <= 1e-2, self._agg_taker_fee
    self.logger.info('%s init pos1 %f' % (self._ticker0, self._get_pos_func()))

  def update_runsize_if_necessary(self, run_size):
    if self._relative_run_size is None or \
        np.isclose(self._relative_run_size, run_size):
      return
    self._relative_run_size = run_size
    overall_run_size = self._raw_run_size * self._relative_run_size
    self._config1 = self._raw_config1.copy()
    self._config1['force_fire'] = True
    max_pos = self._config1['max_pos']
    min_pos = self._config1['min_pos']
    neutral_pos = (max_pos + min_pos) / 2.0
    max_exposure = (max_pos - min_pos) / 2.0 * overall_run_size
    self._config1['max_pos'] = neutral_pos + max_exposure
    self._config1['min_pos'] = neutral_pos - max_exposure
    self._config1['lot_size'] = self._config1['lot_size'] * overall_run_size
    self._equal_theshold = max(self._config1['lot_size'] * 1e-4, 1e-5)
    self._lot_size = self._config1['lot_size']
    self._max_pos = self._config1['max_pos']
    self._min_pos = self._config1['min_pos']
    self.logger.info('Update run_size %f %f %f %f' % (
        self._relative_run_size, self._lot_size, self._config1['min_pos'],
        self._config1['max_pos']))
    self._config1_proto = convert_agg_config(self._product1.symbol, self._config1)
    self._pass_executor.set_config(self._config1_proto)

  def get_statistics(self):
    return self._ma.get_statistics()

  def _update_dq(self, implied_rate):
    self._ma.update(self.get_ts(), implied_rate)

  def _get_ts(self):
    return self._parent._ts

  def get_ts(self):
    return self._get_ts()

  def _is_book_valid(self):
    return self._get_ts() - self._book_ts[0] < self._book_invalid_period

  def _position_equal(self, p1, p2):
    return abs(p1 - p2) < self._equal_theshold

  def get_feed(self, product):
    return self._last_feed.get(product, None)

  def feed_ready(self):
    return self.get_feed(self._product0) is not None and \
        self.get_feed(self._product1) is not None

  def _buy_restrict(self):
    return self._buy_restrict_by_basis or self._buy_restrict_by_book \
        or self._buy_restrict_by_trade

  def _sell_restrict(self):
    return self._sell_restrict_by_basis or self._sell_restrict_by_book \
        or self._sell_restrict_by_trade

  def manage_orders(self):
    bids1, asks1 = self.get_feed(self._product1)
    buy_price = bids1[0][0]
    sell_price = asks1[0][0]
    if self._buy_restrict():
      buy_price = None
    if self._sell_restrict():
      sell_price = None
    cancel_buy = buy_price is None
    cancel_sell = sell_price is None
    self._pass_executor.manage_mm_orders(sell_price,
                                         buy_price,
                                         cancel_all_existing_sell=cancel_sell,
                                         cancel_all_existing_buy=cancel_buy)

  def _book_changed(self, product):
    new_feed = self.get_feed(product)
    last_feed = self._last_book[product]
    if last_feed:
      bids, asks = new_feed
      last_bids, last_asks = last_feed
      self._last_book[product] = new_feed
      for i in range(self._book_change_levels):
        if np.any(bids[i] != last_bids[i]) or np.any(asks[i] != last_asks[i]):
          return True
      return False
    else:
      self._last_book[product] = new_feed
      return True

  def on_book(self, product, book):
    if not self.feed_ready():
      return
    if not self._book_changed(product):
      return

    bids0, asks0 = self.get_feed(self._product0)
    bids1, asks1 = self.get_feed(self._product1)
    self._mid0 = self.get_mid_price(bids0, asks0)
    self._mid1 = self.get_mid_price(bids1, asks1)
    self._implied_scale = self._mid1 / self._mid0
    self._update_dq(self._implied_scale)
    mean, std = self.get_statistics()
    self._buy_restrict_by_basis = \
        self._implied_scale > mean + self._basis_restrict_threshold * std
    self._sell_restrict_by_basis = \
        self._implied_scale < mean - self._basis_restrict_threshold * std
    book_imb = (bids0[0][1] - asks0[0][1]) / (bids0[0][1] + asks0[0][1])
    self._buy_restrict_by_book = book_imb < -self._book_restrict_threshold
    self._sell_restrict_by_book = book_imb > self._book_restrict_threshold

    if self._pass_executor and self.ready():
      self.manage_orders()

  def on_trade(self, product, trade):
    if product == self._product0:
      isbuy = trade.side == TradeSide.TRADE_BUY_SIDE
      if isbuy:
        self._acc_buy_qty += trade.qty
      else:
        self._acc_sell_qty += trade.qty
      if not trade.has_more_trade:
        ts_elapsed = (self.get_ts() - self._last_processed_trade_ts) * 1e-9
        df = math.exp(ts_elapsed * self._lambda)
        self._buy_st = self._buy_st * df + self._acc_buy_qty
        self._sell_st = self._sell_st * df + self._acc_sell_qty
        self._last_processed_trade_ts = self.get_ts()
        self._acc_buy_qty = self._acc_sell_qty = 0
        stat = (self._buy_st - self._sell_st) / (self._buy_st + self._sell_st)
        self._buy_restrict_by_trade = stat < -self._trade_restrict_threshold
        self._sell_restrict_by_trade = stat > self._trade_restrict_threshold
        if self._pass_executor and self.ready():
          self.manage_orders()

  def on_fill(self, fill_price, fill_qty, order):
    self._fillinfo_writer.write(
        FillInfo('Filled',
                 symbol=order.product.symbol,
                 exchange=order.product.exchange,
                 fill_price=fill_price,
                 fill_qty=fill_qty,
                 order_side=order.order_side,
                 timestamp=self._get_ts()))


class BasisPassStrategy(BaseStrategy):
  def __init__(self, strategy_params):
    self.logger = logging.getLogger(__name__)
    self._config = strategy_params['config']
    self._exch0_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch0_subreq']))
    self._exch1_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch1_subreq']))

    self._exchange0 = self._exch0_subreq.exchange
    self._exchange1 = self._exch1_subreq.exchange
    self._og_config = strategy_params['og_config']
    self._balance_checker = None
    account_checker_config = self._get_account_checker_config()
    if account_checker_config is not None:
      self._balance_checker = AccountChecker(account_checker_config)

    self._book_builder_clses = (get_flow_book_builder_cls(self._exch0_subreq),
                                get_flow_book_builder_cls(self._exch1_subreq))

    self._live_env_clses1 = get_live_env_cls(self._exch1_subreq)
    self._live_env1 = None
    self.og1 = None

    self._products0, self._products1 = _gen_arb_products(self._config['common'],
                                                         self._config['products0'],
                                                         self._config['products1'],
                                                         self._book_builder_clses[0].ProductType,
                                                         self._book_builder_clses[1].ProductType)
    self._is_future = self._products1.products[0].product_type == 'Futures'

    exch0_exchange_feed_checker = self._config['common'].get('exch0_exchange_feed_checker', False)
    exch1_exchange_feed_checker = self._config['common'].get('exch1_exchange_feed_checker', False)
    feed_checker_products = []
    if exch0_exchange_feed_checker:
      feed_checker_products += self._products0.products
    if exch1_exchange_feed_checker:
      feed_checker_products += self._products1.products
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)

    multi_feed_subsystem_cls = BaseStrategy.get_multi_feed_subsystem()
    self._feed = multi_feed_subsystem_cls(feed_checker_config=feed_checker_config)
    self._real_time = datetime.datetime.now()
    self._start_time = int(float(self._real_time.strftime('%s.%f')) * 1e9)

    expiry = BaseStrategy.get_config_key(self._products0.products[0])
    self._ref_weekly = 'THIS_WEEK' in expiry
    self._ref_quarter = 'QUARTER' in expiry
    self._run_size = 1.0
    self.logger.info('%s, ref weekly: %s, ref quarter: %s' %
                     (expiry, self._ref_weekly, self._ref_quarter))

    self._ts = None
    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._balance_printer = None
    self._fillinfo_writer = FillInfoWriter('fills.csv', FLAGS.strategy_name)
    self._fill_dumper = FillLogDumper(self._exch1_subreq)

    self._book_levels = FLAGS.book_levels
    self._last_feed = {}
    self._book_map = BookMap(self._products0.products + self._products1.products,
                             logger=self.logger)

    self._reserve_map = self._get_reserve_map()
    self._strategy_request = convert_strategy_request_from_dict(
        self._config['common']['strategy_request'])
    self._reserve_proto = reserve_as_proto(self._exch1_subreq, self._reserve_map)
    self._strategy_logger = StrategyLogProvider(self._strategy_request)

    if FLAGS.enable_telemetry:
      self._latency_recorder = LatencyRecorder('life_of_signal', self._strategy_logger)
    else:
      self._latency_recorder = None

    self._substrategies = {}
    self._usubstrategies = []
    for product0, product1 in zip(self._products0.products, self._products1.products):
      substrategy = BasisPassSubStrategy(self,
                                         product0,
                                         product1,
                                         self._fillinfo_writer,
                                         self._strategy_logger,
                                         self._latency_recorder)
      self._substrategies[product0] = self._substrategies[product1] = substrategy
      self._usubstrategies.append(substrategy)

  def periodic_write_logs(self):
    self.og1.og_logger.write_account_request()
    self._strategy_logger.write_request()
    self._strategy_logger.write_reserve(self._reserve_proto)
    for strat in self._usubstrategies:
      strat.write_executor_config()
    self._ioloop.add_timeout(datetime.timedelta(hours=1), self.periodic_write_logs)

  @property
  def strategy_logger(self):
    return self._strategy_logger

  def _update_prices(self, product, book):
    bids = np.array(book.get_bid_array(self._book_levels))
    asks = np.array(book.get_ask_array(self._book_levels))
    self._last_feed[product] = (bids, asks)

  def on_book(self, product, book):
    self._ts = book.timestamp
    self._update_prices(product, book)
    self._substrategies[product].on_book(product, book)
    self._book_map.set_book(product, book)

  def on_trade(self, product, trade):
    self._ts = trade.timestamp
    self._substrategies[product].on_trade(product, trade)

  def on_fill(self, fill_price, fill_qty, order):
    self._substrategies[order.product].on_fill(fill_price, fill_qty, order)
    self._fill_dumper.dump_fill(fill_price, fill_qty, order)

  def _get_account_checker_config(self):
    if self._config.get('account_checker') is None:
      return None

    config_keys = {
        'reserve',
        'lot_size',
        'stack_size',
        'ratio_min_balance',
        'ratio_max_balance',
        'ratio_min_hold_balance',
    }

    balance_check_config_list = []
    for currency, config in self._config['account_checker'].items():
      config = {**config}
      if 'lot_size' not in config:
        for product in self._config['products']:
          base = SharedProduct.FromStrNativeProduct(product).base.currency
          if base == currency:
            config['lot_size'] = self._config['products'][product]['lot_size']
            break

      assert config_keys == set(config.keys()), config
      reserve = config['reserve']
      config['currency'] = currency
      config['min_balance'] = reserve * config['ratio_min_balance']
      config['max_balance'] = reserve * config['ratio_max_balance']
      config['min_available_balance'] = 0.0
      config['max_available_balance'] = config['lot_size'] * config['stack_size']
      config['min_hold_balance'] = reserve * config['ratio_min_hold_balance']
      config['max_hold_balance'] = reserve * config['ratio_max_balance']
      currency_checker_config = convert_checker_config_from_dict(config)
      balance_check_config_list.append(currency_checker_config)

    return AccountCheckerConfig(each_balance_check_config=balance_check_config_list)

  def signal_handler(self, sig, frame):
    self.logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None and not FLAGS.dry_run:
      ogs = to_list([self.og1])
      self.logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    self._live_env1.start()
    self.start_print_states()

  def subscribe_feeds(self):
    func = self.on_book
    for product in self._products0.products:
      feed_req = FeedSubscriptionRequest.create([product], self._exch0_subreq)
      self._feed.subscribe_book(feed_req, functools.partial(func, product))
      self._feed.subscribe_trade(feed_req, functools.partial(self.on_trade, product))
    for product in self._products1.products:
      feed_req = FeedSubscriptionRequest.create([product], self._exch1_subreq)
      self._feed.subscribe_book(feed_req, functools.partial(func, product))
      # self._feed.subscribe_trade(feed_req, functools.partial(self.on_trade, product))

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    og_params = self._config.get('og_params', {})
    og_params1 = convert_og_config_from_dict(self._exch1_subreq, og_params)

    self._live_env1 = create_live_env_instance(self._live_env_clses1,
                                               self._products1.currencies,
                                               self._products1.products,
                                               self._og_config[0],
                                               og_params=og_params1,
                                               ioloop=ioloop,
                                               queue=queue,
                                               flow_sub=flow_sub,
                                               og_log_callback=self._strategy_logger.write_og_log)
    self._live_env1.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def _get_reserve_map(self):
    reserves = dict()
    for product in self._products1.products:
      key = BaseStrategy.get_config_key(product)
      reserves[product.base.currency] = (self._config['products1'][key]['max_pos']
                                         + self._config['products1'][key]['min_pos']) / 2.0
    return reserves

  def register_balance_printer(self):
    printer_cls = get_pnl_balance_printer_cls(self._exch1_subreq)
    self._balance_printer = printer_cls(self._products1.products,
                                        self._book_map,
                                        self._reserve_map,
                                        logger=self.logger)
    self._balance_printer.register_order_gateway(self.og1)
    PeriodicCallback(self.dump_pnl_balance, 30 * 1000).start()

  def start_print_states(self):
    for strat in self._usubstrategies:
      # TODO print states of strategy.
      break
    self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.start_print_states)

  def dump_pnl_balance(self, periodic=False):
    if self._balance_printer is not None:
      self._balance_printer.print_pnl_balance()
      pnl_balances = self._balance_printer.calculate_pnl_balance()
      if pnl_balances is not None:
        for pnl_balance in pnl_balances:
          self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance,
                                                  quote=None,
                                                  subreq_id=None)
    self._update_runsize_if_necessary()

  def _update_runsize_if_necessary(self):
    self._real_time = datetime.datetime.now()
    run_size = BaseStrategy.check_rollover(self._ref_weekly, self._ref_quarter, self._real_time)
    schedules = BaseStrategy.read_maintenance_as_json(FLAGS.maintenance_file,
                                                      [self._exchange0, self._exchange1])
    for p0, p1 in zip(self._products0.products, self._products1.products):
      run_size2 = min(BaseStrategy.check_maintenance(schedules, p0, self._real_time),
                      BaseStrategy.check_maintenance(schedules, p1, self._real_time))
      run_size2 = min(run_size, run_size2)
      self._substrategies[p1].update_runsize_if_necessary(run_size2)
    if schedules:
      self.logger.info('Schedules count %s' % len(schedules))

  def register_order_gateway(self, name, order_gateway):
    self.logger.info('OG initialized: %s', name)
    if order_gateway.exchange == self._exchange1:
      self.og1 = order_gateway
    # cancel working orders first after OGs are ready.
    if not self._order_canceled_after_og_ready:
      order_gateway.cancel_all()
      self.logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True
    self._exit_executor = SimpleCancelAllExecutor()

    if self.og1:
      self.register_balance_printer()
      self.og1.register_fill_callback(self.on_fill)
      for prod0, prod1 in zip(self._products0.products, self._products1.products):
        key = BaseStrategy.get_config_key(prod1)
        config1 = self._config['products1'][key].copy()
        config1.update(self._config.get('pass_executor_configs', {}))
        self._substrategies[prod1].init_executors(self._exch1_subreq, config1, self.og1)
      self.periodic_write_logs()


def get_param_set(og_configs, config_filename):
  og_config = og_configs.split(',')
  config = load_clean_json(config_filename)
  return {
      'og_config': og_config,
      'config': config,
  }


def main(argv):
  assert FLAGS.strategy_name

  strategy_params = get_param_set(og_configs=FLAGS.og_configs,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('basisunhedged')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()

  strategy = BasisPassStrategy(strategy_params)
  strategy.subscribe_feeds()
  strategy.subscribe_live_envs(ioloop, queue, flow_sub)
  strategy.start(ioloop, queue, flow_sub)

  # exit_after = datetime.timedelta(days=1)
  # ioloop.add_timeout(exit_after, strategy.on_exit)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  ioloop.start()


if __name__ == '__main__':
  flags.DEFINE_string('og_configs', '', '')

  flags.DEFINE_string('config_filename', '', '')

  flags.DEFINE_boolean('dry_run', False, '')

  flags.DEFINE_float('book_invalid_sec', 3.5, '')

  flags.DEFINE_boolean('update_pos_with_balance', False, '')

  flags.DEFINE_float('update_pos_with_balance_sec', 4, '')

  flags.DEFINE_float('update_working_qty_sec', 3, '')

  flags.DEFINE_integer('book_levels', 5, '')

  flags.DEFINE_integer('book_change_levels', 3, '')

  flags.DEFINE_string('strategy_name', '', '')

  flags.DEFINE_float('run_size', 1.0, '')

  flags.DEFINE_string('maintenance_file', 'coin_deploy/basis_strat/schedule/maintenance.json', '')

  flags.DEFINE_bool(
      'enable_telemetry',
      False,
      'enable telemetry',
  )

  flags.DEFINE_integer('window_period_sec', 300, '')
  flags.DEFINE_float('basis_restrict_threshold', 2.0, '')
  flags.DEFINE_float('book_imbalance_restrict_threshold', 0.7, '')
  flags.DEFINE_float('trade_imbalance_restrict_threshold', 0.4, '')
  flags.DEFINE_float('discount_factor', -0.25, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
